from __future__ import division, print_function

import os
# Ignore warnings
import warnings

import cv2
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from genericpath import samestat
from matplotlib import cm, image
from matplotlib.cm import ScalarMappable
from skimage import io, transform
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms, utils

from data_pipeline import Data_Pipeline
from transforms import classfifation_transform

warnings.filterwarnings("ignore")

mnist = Data_Pipeline(label_file='dataset/MNIST/mnist_datatrain.txt',
                                    root_dir='dataset/MNIST/mnist_data/train',
                                    transform=classfifation_transform)

for i in range(len(mnist)):

    sample = mnist[i]

    # print(i, sample['image'].shape, type(sample['label']))

    ax = plt.subplot(1, 4, i + 1)
    plt.tight_layout()
    ax.set_title('label: {}'.format(sample['label']))
    ax.axis('off')
    plt.imshow(sample['image'])

    if i == 3:
        plt.show()
        break
